import numpy as np
import matplotlib.pyplot as plt
from init_control import get_control_parameters
from pid_controller import PIDController, CascadePIDController
from uav_dynamics import uav_dynamics
from scipy.spatial.transform import Rotation as R

def motor_mixer(roll, pitch, yaw, thrust):
    """
    Control allocation. The quadrotor type is X-configuration,
    and the airframe is as follows:
    3↓   1↑
      \ /
      / \
    2↑   4↓

    Input:
      roll, pitch, yaw: attitude controller output.
      thrust: total thrust.

    Output:
      M1, M2, M3, M4: motor speed commands.
    """
    idle_PWM = 1000
    scale = 1000

    M1 = (thrust - roll + pitch + yaw) * scale + idle_PWM
    M2 = (thrust + roll - pitch + yaw) * scale + idle_PWM
    M3 = (thrust + roll + pitch - yaw) * scale + idle_PWM
    M4 = (thrust - roll - pitch - yaw) * scale + idle_PWM

    return M1, M2, M3, M4


def desired_euler_angles(eax, eay):
    """
    Convert the desired horizontal acceleration to the desired Euler angle

    Input:
      eax, eay: desired horizontal acceleration

    Output:
      phi, theta: desired roll and pitch
    """
    g = 9.8  # 重力加速度 (m/s^2)

    phi = eay / g
    theta = -eax / g

    return phi, theta


def pwm_to_rpm(pwm):
    """
    Convert PWM signal to RPM. This is a simple linear approximation.
    You might need to adjust the conversion based on actual hardware data.

    Input:
      pwm: PWM signal value

    Output:
      rpm: motor speed in RPM
    """
    temp_pwm= (pwm - 1000) / 1000
    print('temp_pwm:',temp_pwm)
    if temp_pwm>1:
        return 1
    elif temp_pwm<-1:
        return -1
    else:
        return temp_pwm
    # pwms_in_motor=(pwm - 1000) / 1000
    # gain = pwms_in_motor*ModelParam_motorCr+ModelParam_motorWb
    # return (gain/ModelParam_motorT)*math.exp(-)


def run_uav_simulation():
    params = get_control_parameters()

    # 初始化PID控制器
    # roll_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0)
    # roll_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
    #                                      params['Kd_RP_AngleRate'],
    #                                      integrator_max=params['Saturation_I_RP_Max'],
    #                                      integrator_min=params['Saturation_I_RP_Min'])
    # roll_controller = CascadePIDController(roll_angle_controller, roll_rate_controller)
    #
    # pitch_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0)
    # pitch_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
    #                                       params['Kd_RP_AngleRate'],
    #                                       integrator_max=params['Saturation_I_RP_Max'],
    #                                       integrator_min=params['Saturation_I_RP_Min'])
    # pitch_controller = CascadePIDController(pitch_angle_controller, pitch_rate_controller)
    #
    # yaw_rate_controller = PIDController(params['Kp_YAW_AngleRate'], params['Ki_YAW_AngleRate'],
    #                                     params['Kd_YAW_AngleRate'],
    #                                     integrator_max=params['Saturation_I_Y_Max'],
    #                                     integrator_min=params['Saturation_I_Y_Min'])
    roll_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0,integrator_max= 1,
                                         integrator_min=-1)
    roll_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
                                         params['Kd_RP_AngleRate'],
                                         integrator_max= 1,
                                         integrator_min=-1)
    roll_controller = CascadePIDController(roll_angle_controller, roll_rate_controller)

    pitch_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0,integrator_max= 1,
                                         integrator_min=-1)
    pitch_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
                                          params['Kd_RP_AngleRate'],
                                          integrator_max=1,
                                          integrator_min=-1)
    pitch_controller = CascadePIDController(pitch_angle_controller, pitch_rate_controller)

    yaw_rate_controller = PIDController(params['Kp_YAW_AngleRate'], params['Ki_YAW_AngleRate'],
                                        params['Kd_YAW_AngleRate'],
                                        integrator_max=1,
                                        integrator_min=-1)

    # 添加高度控制PID
    height_controller = PIDController(0.6, 0.3, 0.05)  # 调整这些值以获得稳定的高度控制

    # 初始化仿真状态
    position = params['ModelInit_PosE'].astype(np.float64)
    velocity_body = params['ModelInit_VelB'].astype(np.float64)
    velocity_world = 0
    angle = params['ModelInit_AngEuler'].astype(np.float64)
    rate = params['ModelInit_RateB'].astype(np.float64)
    w = np.zeros(4, dtype=np.float64)  # 电机转速
    mass = params['ModelParam_uavMass']
    J = params['ModelParam_uavJ']

    dt = 0.001  # 时间步长，单位为秒
    simulation_time = 20  # 仿真总时间，单位为秒
    time_steps = int(simulation_time / dt)

    # 定义起飞阶段和正方形轨迹的四个角点
    target_altitude = -5  # 目标高度
    takeoff_time = 5  # 起飞时间，单位为秒
    takeoff_steps = int(takeoff_time / dt)

    square_corners = [
        np.array([0, 0, target_altitude]),  # 起点
        np.array([5, 0, target_altitude]),
        np.array([5, 5, target_altitude]),
        np.array([0, 5, target_altitude]),
        np.array([0, 0, target_altitude])  # 回到起点
    ]

    current_corner_index = 0
    next_corner_time = (time_steps - takeoff_steps) // len(square_corners)

    # 数据记录
    position_history = []

    for step in range(time_steps):
        if step < takeoff_steps:
            # 起飞阶段，垂直上升到目标高度
            setpoint = np.array([0, 0, target_altitude])
        else:
            # 正方形轨迹阶段
            # setpoint = (square_corners[current_corner_index]-square_corners[current_corner_index-1])/((time_steps-takeoff_steps)/5)
            setpoint = square_corners[current_corner_index]
        # 位置控制器

        # 计算期望的水平加速度
        eax = (setpoint[0] - position[0]) / (time_steps - takeoff_steps)
        eay = (setpoint[1] - position[1]) / (time_steps - takeoff_steps)#[0]

        # 期望的欧拉角
        phi, theta = desired_euler_angles(eax, eay)

        # 更新姿态控制器输出
        roll_output = roll_controller.update(outer_setpoint=phi, outer_measurement=angle[0], inner_measurement=rate[0],
                                             dt=dt,)
        pitch_output = pitch_controller.update(outer_setpoint=theta, outer_measurement=angle[1],
                                               inner_measurement=rate[1], dt=dt)
        yaw_output = yaw_rate_controller.update(setpoint=0, measurement=rate[2], dt=dt)

        # 高度控制
        altitude_output = height_controller.update(setpoint[2], position[2], dt)

        # 计算推力，考虑重力补偿
        thrust = altitude_output * params['ModelParam_uavMass'] * params['ModelParam_envGravityAcc']

        # 计算电机转速
        M1, M2, M3, M4 = motor_mixer(roll_output, pitch_output, yaw_output, thrust)
        # 将PWM信号转换为转速（假设线性关系）
        w = np.array([pwm_to_rpm(M1), pwm_to_rpm(M2), pwm_to_rpm(M3), pwm_to_rpm(M4)])
        print(setpoint)


        # 更新无人机动力学模型
        # position, velocity_world, angle, rate = uav_dynamics(position, velocity_body, angle, rate, w, dt, params)
        total_force,total_torque = uav_dynamics(position, velocity_body, angle, rate, w, dt, params)
        # 合力计算

        acceleration = total_force / mass
        velocity_world += acceleration * dt
        position += velocity_world * dt
        position_history.append(position.copy())

        angular_acceleration = np.linalg.inv(J).dot(total_torque)
        rate += angular_acceleration * dt
        angle += rate * dt

        r = R.from_euler('xyz', angle, degrees=False)
        velocity_body = r.inv().apply(velocity_world)
        # 打印当前仿真状态（仅供示例）
        print(f"Time: {step * dt:.2f}s, Position: {position}, Velocity: {velocity_world}, Angle: {angle}, Rate: {rate}")

        # 更新目标点索引
        if step >= takeoff_steps and (step - takeoff_steps) % next_corner_time == 0 and step != takeoff_steps:
            current_corner_index = (current_corner_index + 1) % len(square_corners)

    # 转换为numpy数组
    position_history = np.array(position_history)

    # 绘制三维轨迹图
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(position_history[:, 0], position_history[:, 1], position_history[:, 2], label='UAV Trajectory')
    ax.set_xlabel('X Position (m)')
    ax.set_ylabel('Y Position (m)')
    ax.set_zlabel('Z Position (m)')
    ax.set_title('3D Trajectory of UAV Takeoff and Square Path')
    ax.legend()
    plt.show()
# def run_uav_simulation():
#     params = get_control_parameters()
#
#     # 初始化PID控制器
#     roll_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0)
#     roll_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
#                                          params['Kd_RP_AngleRate'],
#                                          integrator_max=params['Saturation_I_RP_Max'],
#                                          integrator_min=params['Saturation_I_RP_Min'])
#     roll_controller = CascadePIDController(roll_angle_controller, roll_rate_controller)
#
#     pitch_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0)
#     pitch_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
#                                           params['Kd_RP_AngleRate'],
#                                           integrator_max=params['Saturation_I_RP_Max'],
#                                           integrator_min=params['Saturation_I_RP_Min'])
#     pitch_controller = CascadePIDController(pitch_angle_controller, pitch_rate_controller)
#
#     yaw_rate_controller = PIDController(params['Kp_YAW_AngleRate'], params['Ki_YAW_AngleRate'],
#                                         params['Kd_YAW_AngleRate'],
#                                         integrator_max=params['Saturation_I_Y_Max'],
#                                         integrator_min=params['Saturation_I_Y_Min'])
#
#     # 添加高度控制PID
#     height_controller = PIDController(0.6, 0.3, 0.05)  # 调整这些值以获得稳定的高度控制
#
#     # 初始化仿真状态
#     position = params['ModelInit_PosE'].astype(np.float64)
#     velocity = params['ModelInit_VelB'].astype(np.float64)
#     angle = params['ModelInit_AngEuler'].astype(np.float64)
#     rate = params['ModelInit_RateB'].astype(np.float64)
#     w = np.zeros(4, dtype=np.float64)  # 电机转速
#
#     dt = 0.01  # 时间步长，单位为秒
#     simulation_time = 60  # 仿真总时间，单位为秒
#     time_steps = int(simulation_time / dt)
#
#     # 定义起飞阶段和正方形轨迹的四个角点
#     target_altitude = -5  # 目标高度
#     takeoff_time = 10  # 起飞时间，单位为秒
#     takeoff_steps = int(takeoff_time / dt)
#
#     square_corners = [
#         np.array([0, 0, target_altitude]),  # 起点
#         np.array([5, 0, target_altitude]),
#         np.array([5, 5, target_altitude]),
#         np.array([0, 5, target_altitude]),
#         np.array([0, 0, target_altitude])  # 回到起点
#     ]
#
#     current_corner_index = 0
#     next_corner_time = (time_steps - takeoff_steps) // len(square_corners)
#
#     # 数据记录
#     position_history = []
#
#     for step in range(time_steps):
#         if step < takeoff_steps:
#             # 起飞阶段，垂直上升到目标高度
#             setpoint = np.array([0, 0, target_altitude])
#         else:
#             # 正方形轨迹阶段
#             setpoint = square_corners[current_corner_index]
#
#         # 计算位置误差
#         eax = setpoint[0] - position[0]
#         eay = setpoint[1] - position[1]
#
#         # 期望的欧拉角
#         phi, theta = desired_euler_angles(eax, eay)
#
#         # 更新控制器输出
#         roll_output = roll_controller.update(outer_setpoint=phi, outer_measurement=angle[0], inner_measurement=rate[0],
#                                              dt=dt)
#         pitch_output = pitch_controller.update(outer_setpoint=theta, outer_measurement=angle[1],
#                                                inner_measurement=rate[1], dt=dt)
#         yaw_output = yaw_rate_controller.update(setpoint=0, measurement=rate[2], dt=dt)
#
#         # 高度控制
#         altitude_output = height_controller.update(setpoint[2], position[2], dt)
#
#         # 计算推力，考虑重力补偿
#         thrust = (altitude_output + 1) * params['ModelParam_uavMass'] * params['ModelParam_envGravityAcc']
#
#         # 计算电机转速
#         M1, M2, M3, M4 = motor_mixer(roll_output, pitch_output, yaw_output, thrust)
#
#         # 将PWM信号转换为转速（假设线性关系）
#         w = np.array([M1, M2, M3, M4])
#
#         # 更新无人机动力学模型
#         position, velocity, angle, rate = uav_dynamics(position, velocity, angle, rate, w, dt, params)
#
#         # 记录位置
#         position_history.append(position.copy())
#
#         # 打印当前仿真状态（仅供示例）
#         print(f"Time: {step * dt:.2f}s, Position: {position}, Velocity: {velocity}, Angle: {angle}, Rate: {rate}")
#
#         # 更新目标点索引
#         if step >= takeoff_steps and (step - takeoff_steps) % next_corner_time == 0 and step != takeoff_steps:
#             current_corner_index = (current_corner_index + 1) % len(square_corners)
#
#     # 转换为numpy数组
#     position_history = np.array(position_history)
#
#     # 绘制三维轨迹图
#     fig = plt.figure()
#     ax = fig.add_subplot(111, projection='3d')
#     ax.plot(position_history[:, 0], position_history[:, 1], position_history[:, 2], label='UAV Trajectory')
#     ax.set_xlabel('X Position (m)')
#     ax.set_ylabel('Y Position (m)')
#     ax.set_zlabel('Z Position (m)')
#     ax.set_title('UAV Trajectory')
#     ax.legend()
#     plt.show()


if __name__ == "__main__":
    run_uav_simulation()
